package com.jscloud.bigdata.flink.datastreamdemo;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

/**
 * - 如果Flink自带的分区方式还没法满足我们的需求，我们还可以自定义分区策略来实现数据的分区
 * - 需求
 * - 自定义分区策略，实现不同分区的数据发送到不同分区里面去进行处理，将包含hello的字符串发送到一个分区里面去，
 * 其他的发送到另外一个分区里面去
 */
public class FlinkCustomerPartitionByJava {
        public static void main(String[] args) throws Exception {
                StreamExecutionEnvironment executionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment();
                DataStreamSource<String> sourceStream = executionEnvironment.fromElements("hello world", "hello spark", "hudi hive", "test result");

                //自定义分区策略
                DataStream<String> customPartition = sourceStream.partitionCustom(new MyPartitioner(), new KeySelector<String, String>() {
                        @Override
                        public String getKey(String value) throws Exception {
                                return value;
                        }
                });


                SingleOutputStreamOperator<String> mapStream = customPartition.map(new MapFunction<String, String>() {
                        @Override
                        public String map(String value) throws Exception {
                                System.out.println("当前线程ID为" + Thread.currentThread().getId());
                                return value;
                        }
                });

                mapStream.print();
                executionEnvironment.execute();
        }
}

class MyPartitioner implements Partitioner<String> { //自定义一定要实现 implements Partitioner
        @Override
        public int partition(String key, int numPartitions) {
                System.out.println(numPartitions);
                if (key.contains("hello")) {
                        return 0;
                }
                return 1;
        }
}